#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
'''
@File         :path_vir_leader.py
@Description  :位置环控制，主要用于领机的offb控制
@Time         :2024/05/30 16:54:00
@Author       :Lin Yuheng
@Version      :1.0
'''

import rospy
from geometry_msgs.msg import Pose, PoseStamped, TwistStamped
from outdoorground.msg import FWstates, Leaderstates
import math
import numpy as np

STATE_UAV = 101
CMD_UAV = 102
ARM = 1
TAKEOFF = 2
LAND = 3
POS = 4
VEL = 5
ACC = 6
POS_VEL = 7
POS_ACC = 8
VEL_ACC = 9
POS_VEL_ACC = 10
ATT = 11

class Controller:
    def __init__(self):
        self.plane_id = 0
        self.vehicle_id = "plane_1"
        self.leader_id = "plane_0"
        self.px = 0
        self.py = 0
        self.pz = 0
        self.vx, self.vy, self.vz = 0, 0, 0
        self.roll, self.pitch, self.yaw = 0, 0, 0
        self.phi_dot, self.theta_dot, self.psi_dot = 0, 0, 0
        self.send_flag = False
        
        self.vir_leader_state = Leaderstates()

        # subscribers
        self.local_pose_sub = rospy.Subscriber(
            self.vehicle_id + "/fixed_wing_control/fw_state",
            FWstates,
            callback=self.local_pose_callback,
        )

        # publishers
        # 发布虚拟领导者的状态信息
        self.vir_leader_state_pub = rospy.Publisher(
            self.leader_id + "/fixed_wing_control/fw_state",
            Leaderstates,
            queue_size=1
        )

    def local_pose_callback(self, msg: FWstates):
        self.px = msg.ned_pos_x
        self.py = msg.ned_pos_y
        self.pz = msg.ned_pos_z

        self.vx = msg.ned_vel_x
        self.vy = msg.ned_vel_y
        self.vz = msg.ned_vel_z
        
    def interpolate_linear(self, p1, p2, t):
        return (p1[0] + (p2[0] - p1[0]) * t, p1[1] + (p2[1] - p1[1]) * t)

    def interpolate_arc(self, center, radius: float, start_angle: float, end_angle: float, t: float):
        start_angle = math.pi / 2 - start_angle
        end_angle = math.pi / 2 - end_angle
        angle = start_angle + (end_angle - start_angle) * t
        return (center[0] + radius * math.cos(angle), center[1] + radius * math.sin(angle))
    
    # 计算当前位置到点的距离
    def distance_to_setpoint(self, setpoint):
        return math.sqrt((self.px - setpoint[0]) ** 2 + (self.py - setpoint[1]) ** 2)

    def start(self):
        print("start path_vir_leader_state send")
        rospy.init_node(
            "path_vir_leader_node"
        )
        rate = rospy.Rate(20)
        
        waypoint_start = (0, -100)
        waypoint = waypoint_start
        reach_start = False
        
        upward = 50
        speed_cruise = 15.5
        
        # 轨迹的关键点 enu坐标系
        A = (0.0, 150.0)
        B = (200.0, 150.0)
        C = (200.0, -100.0)
        D = (0.0, -100.0)
        
        
        # 半圆的中心和半径
        center1 = (100.0, 150)  # 半圆 1 的中心
        radius1 = 100
        center2 = (100.0, -100)  # 半圆 2 的中心
        radius2 = 100
        
        # 分段长度
        len1 = 250 # 第一段直线长度
        len2 = math.pi * radius1  # 第二段半圆长度
        len3 = 250  # 第三段直线长度
        len4 = math.pi * radius2  # 第四段半圆长度
        length_total = len1 + len2 + len3 + len4
        
        time_total = length_total / speed_cruise

        # 计算每段轨迹的参数范围
        t1 = len1 / length_total * time_total
        t2 = t1 + len2 / length_total * time_total
        t3 = t2 + len3 / length_total * time_total
        
        wp_prev = waypoint

        while not rospy.is_shutdown():
            current_time = rospy.Time.now()
            dist = self.distance_to_setpoint(waypoint)
            if not reach_start:
                cmd_type = POS
                v_sp_x = 0
                v_sp_y = speed_cruise
                if dist < 30:
                    reach_start = True
                    t_init = current_time.to_sec()
            else:
                cmd_type = ATT
                t = current_time.to_sec() - t_init
                
                while t >= time_total:
                    t = t - time_total
                if t <= t1:
                    # 第一段直线
                    waypoint = self.interpolate_linear(D, A, t / t1)
                elif t <= t2:
                    # 第二段半圆
                    waypoint = self.interpolate_arc(center1, radius1, -math.pi/2, math.pi/2, (t - t1) / (t2 - t1))
                elif t <= t3:
                    # 第三段直线
                    waypoint = self.interpolate_linear(B, C, (t - t2) / (t3 - t2))
                else:
                    # 第四段半圆
                    waypoint = self.interpolate_arc(center2, radius2, math.pi/2, 3 * math.pi/2, (t - t3) / (time_total - t3))
                    
                v_sp = np.sqrt((waypoint[0] - wp_prev[0]) ** 2 + (waypoint[1] - wp_prev[1]) ** 2) / 0.02
                v_sp_x = (waypoint[0] - wp_prev[0]) / 0.02
                v_sp_y = (waypoint[1] - wp_prev[1]) / 0.02
                wp_prev = waypoint
                # print("v_sp",v_sp)
            
            self.vir_leader_state.cmd_type = cmd_type
            self.vir_leader_state.airspeed = speed_cruise
            self.vir_leader_state.ned_pos_x = waypoint[1]
            self.vir_leader_state.ned_pos_y = waypoint[0]
            self.vir_leader_state.ned_pos_z = - upward
            self.vir_leader_state.ned_vel_x = v_sp_y
            self.vir_leader_state.ned_vel_y = v_sp_x
            self.vir_leader_state.ned_vel_z = 0
            self.vir_leader_state.global_vel_x = v_sp_y
            self.vir_leader_state.global_vel_y = v_sp_x
            self.vir_leader_state.global_vel_z = 0

            try:
                self.vir_leader_state_pub.publish(self.vir_leader_state)
            except rospy.ROSException as e:
                pass
            rate.sleep()
            
            if cmd_type == POS:
                cmd_str = "POS"
            elif cmd_type == ATT:
                cmd_str = "ATT"
            rospy.loginfo(
                "cmd_type: %s, wp_sp: (%f, %f); fw_pos: (%f, %f); distance:%f",cmd_str, waypoint[1], waypoint[0], self.px, self.py, dist
            )


if __name__ == "__main__":
    Con = Controller()
    try:
        Con.start()
    except rospy.ROSInterruptException:
        pass
